import numpy as np
import cv2,time


def label_for_patch(patch_alpha,patch_image,threshold=0.99):
    '''
    Args:
          patch_alpha: ground-truth alpha matte patch [h,c,1]
          patch_image: input colorful image patch [h,c,3]
    '''
    image_height,image_width,_ = patch_image.shape
    num_bg = 0
    num_fg = 0
    for i in range (image_height):
        for j in range (image_width):
            if patch_alpha[i,j,0]==0:
               num_bg+=1
            elif patch_alpha[i,j,0]==255:
               num_fg+=1
    all_pixels = image_height*image_width
    bg_ratio = num_bg/all_pixels
    fg_ratio = num_fg/all_pixels
    if bg_ratio > threshold:
       patch_image[:,:]=(0,0,0)
    elif fg_ratio>threshold:
         patch_image[:,:]=(255,255,255)
    else:
         patch_image[:,:]=(128,128,128)
    return patch_image.copy()


def label_for_image(trimap,image,alpha,loc,win_size):
    '''
    Args:
          trimap: accumulate trimap
          image: colorful image input
          alpha: ground-truth alpha matte
          loc: location of label
    '''
    image_height,image_width,_ = image.shape
    leftup=np.array([max(0,(loc[0] - win_size[0]/2)),max(0,(loc[1] - win_size[1]/2))],np.int32)
    rightdown =np.array([min(image_height,(loc[0] + win_size[0]/2)),min(image_width,(loc[1] + win_size[1]/2))],np.int32)
    

    patch_image = image[leftup[0]:rightdown[0],leftup[1]:rightdown[1]]
    patch_alpha = alpha[leftup[0]:rightdown[0],leftup[1]:rightdown[1]] 
 
    
    patch_trimap = label_for_patch(patch_alpha,patch_image)
    trimap[leftup[0]:rightdown[0],leftup[1]:rightdown[1]]=patch_trimap
    return trimap


def user_interaction(image,loc,win_size):
    def draw_call(event,x,y,flags,param):
        
        if event == cv2.EVENT_LBUTTONDOWN:
            if cmode == True:        #Foreground
               cv2.rectangle(image_copy,leftup,rightdown,(255,255,255),-1)
            else:                    #Background
               cv2.rectangle(image_copy,leftup,rightdown,(0,0,0),-1)
    
    def shadow_image(image_copy,loc,win_size):
        image_copy_copy = image_copy.copy()
        cv2.rectangle(image_copy_copy,leftup,rightdown,(0,0,255),2)
        return image_copy_copy
    image_copy = image.copy()
    image_height,image_width,_ = image.shape
    shape = np.array((image_height,image_width))
    loc = np.array((loc*shape/2+shape/2),np.int32) # loc from [-1,1] to the real location coordinates
    leftup=tuple(np.array([max(0,(loc[0] - win_size[0]/2)),max(0,(loc[1] - win_size[1]/2))],np.int32))
    rightdown =tuple(np.array([min(image_height,(loc[0] + win_size[0]/2)),min(image_width,(loc[1] + win_size[1]/2))],np.int32))
    cmode = False
    cv2.namedWindow('Image')
    cv2.setMouseCallback('Image',draw_call)
    while(True):
         image_c_c_c = shadow_image(image_copy,loc,win_size)
         cv2.imshow('Image',image_c_c_c)
         k = cv2.waitKey(1)&0xFF
         if k == ord('f'):     #Foreground
            cmode = True
         elif k == ord('b'):   #Background
              cmode = False
         elif k == ord('n'):   #SAVE
              image = image_copy
         elif k==27:      #ESC
              break
         time.sleep(0.05)
    cv2.destroyAllWindows()
    return image
    



def label_for_batches(trimaps,alphas,locs,win_size,mode='Training'):
    '''
    Args:
         win_size: an array defining the size of labeled windows
         mode: Training or Testing, when training, use label_for_image(); for testing,use user_interaction()
    '''
    batch_num,height,width,_ = trimaps.shape
    images = trimaps
    for i in range(batch_num):
        loc = locs[i]
        image = images[i]
        alpha = alphas[i]
        trimap = trimaps[i]
        if mode == 'Training':
           shape = np.array((height,width))
           loc = np.array((loc*shape/2+shape/2),np.int32) # loc from [-1,1] to the real location coordinates
           labeled_trimap = label_for_image(trimap,image,alpha,loc,win_size)
        else:
            labeled_trimap = user_interaction(image,loc,win_size)
        trimaps[i] = labeled_trimap
    return trimaps


if __name__=='__main__':
   image = cv2.imread('./data/image.png')
   alpha = cv2.imread('./data/alpha.png')
   loc=[0,-0.3]
   win_size = [150,206]
   trimaps=label_for_batches(np.array([image]),np.array([alpha]),[loc],win_size,mode='Testing')
   trimap = trimaps[0]
   #trimap=user_interaction(image,loc,win_size)
   cv2.imwrite('./data/trimap.png',trimap)

